import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs

class KMeans:
    def __init__(self, k=3, max_iters=100, tol=1e-4):
        self.k = k
        self.max_iters = max_iters
        self.tol = tol
        self.centroids = None
        self.labels = None
        
    def fit(self, X):
        """训练KMeans模型"""
        n_samples, n_features = X.shape
        
        # 随机初始化中心点
        indices = np.random.choice(n_samples, self.k, replace=False)
        self.centroids = X[indices]
        
        for _ in range(self.max_iters):
            # 分配样本到最近的中心点
            distances = self._compute_distances(X)
            self.labels = np.argmin(distances, axis=1)
            
            # 更新中心点
            new_centroids = np.array([X[self.labels == i].mean(axis=0) for i in range(self.k)])
            
            # 检查收敛
            if np.linalg.norm(new_centroids - self.centroids) < self.tol:
                break
                
            self.centroids = new_centroids
            
        return self
    
    def _compute_distances(self, X):
        """计算样本到所有中心点的距离"""
        distances = np.zeros((X.shape[0], self.k))
        for i, centroid in enumerate(self.centroids):
            distances[:, i] = np.linalg.norm(X - centroid, axis=1)
        return distances
    
    def predict(self, X):
        """预测样本所属聚类"""
        distances = self._compute_distances(X)
        return np.argmin(distances, axis=1)

# 测试示例
if __name__ == "__main__":
    # 生成测试数据
    X, _ = make_blobs(n_samples=300, centers=3, n_features=2, random_state=42)
    
    # 训练KMeans
    kmeans = KMeans(k=3)
    kmeans.fit(X)
    labels = kmeans.labels
    
    # 可视化结果
    plt.figure(figsize=(10, 4))
    
    plt.subplot(121)
    plt.scatter(X[:, 0], X[:, 1], c='gray', alpha=0.6)
    plt.title("原始数据")
    
    plt.subplot(122)
    for i in range(3):
        cluster_points = X[labels == i]
        plt.scatter(cluster_points[:, 0], cluster_points[:, 1], label=f'Cluster {i}')
    plt.scatter(kmeans.centroids[:, 0], kmeans.centroids[:, 1], 
                marker='x', s=200, linewidths=3, color='black', label='Centroids')
    plt.title("KMeans聚类结果")
    plt.legend()
    
    plt.tight_layout()
    plt.show()
    
    print("聚类中心坐标:")
    print(kmeans.centroids)